
'''
init test dataset
'''

import os
import numpy as np
import random
import matplotlib.pyplot as plt
import librosa
import librosa.display
from tqdm import tqdm

class TestDatasetTools:

    def __init__(self, raw_path, dataset_path):
        
        self.raw_path = raw_path
        self.dataset_path = dataset_path
        self.audio_length = 100000

        if not os.path.exists(self.dataset_path):
            os.mkdir(self.dataset_path)
        

    def convert_audio_to_image(self, audio_path, image_dir):

        raw_data, raw_framesize = librosa.load(audio_path, sr=None, mono=False)

        raw_length = len(raw_data)
        if raw_length <= self.audio_length:
            # need to add a new part in the end
            
            raw_data_final = raw_data[100:-100]
            repeat_num = int(self.audio_length / raw_length) + 1
            for times in range(repeat_num+1):
                raw_data_final = np.hstack((raw_data_final, raw_data[100:-100]))

            raw_length = len(raw_data_final)

            max_start_idx = raw_length - self.audio_length

            name_idx = 0
            
            for img_idx in range(10):
                start_idx = random.randint(0, max_start_idx)
                clipped_data = raw_data_final[start_idx: start_idx + self.audio_length]

                data = clipped_data * 1.0 / clipped_data.max()
                framelength = 0.025
                framesize = int(framelength * raw_framesize)

                # get feature
                mel_spect = librosa.feature.melspectrogram(data, sr=raw_framesize, n_fft=framesize)
                mel_spect = librosa.power_to_db(mel_spect, ref=np.max)

                # draw
                image_path = os.path.join(image_dir, \
                                os.path.basename(audio_path)[:-4] + '_' + str(name_idx) + '.jpg')
                plt.axis('off')
                librosa.display.specshow(mel_spect, sr=raw_framesize, x_axis='time', y_axis='mel')
                plt.savefig(image_path, transparent=False, bbox_inches='tight', pad_inches=0)
                plt.close()
                name_idx += 1

        else:
            # clip

            max_start_idx = raw_length - self.audio_length

            name_idx = 0
            
            for img_idx in range(10):
                start_idx = random.randint(0, max_start_idx)
                clipped_data = raw_data[start_idx: start_idx + self.audio_length]

                data = clipped_data * 1.0 / clipped_data.max()
                framelength = 0.025
                framesize = int(framelength * raw_framesize)

                # get feature
                mel_spect = librosa.feature.melspectrogram(data, sr=raw_framesize, n_fft=framesize)
                mel_spect = librosa.power_to_db(mel_spect, ref=np.max)

                # draw
                image_path = os.path.join(image_dir, \
                                os.path.basename(audio_path)[:-4] + '_' + str(name_idx) + '.jpg')
                plt.axis('off')
                librosa.display.specshow(mel_spect, sr=raw_framesize, x_axis='time', y_axis='mel')
                plt.savefig(image_path, transparent=False, bbox_inches='tight', pad_inches=0)
                plt.close()
                name_idx += 1

    def check_path(self):
        
        test_root_path = self.dataset_path
        if not os.path.exists(test_root_path):
            os.mkdir(test_root_path)
    

    def start_convert(self):

        self.check_path()
        audio_list = os.listdir(self.raw_path)
        
        for idx in tqdm(range(len(audio_list))):
            audio_path = audio_list[idx]
            audio_file = os.path.join(self.raw_path, audio_path)
            image_dir = os.path.join(self.dataset_path, os.path.basename(audio_file)[:-4])
            if not os.path.exists(image_dir):
                os.mkdir(image_dir)
            self.convert_audio_to_image(audio_file, image_dir)


# if __name__ == '__main__':
#     dataset_tools = TestDatasetTools('../../../dataset/test_offline/task2', './')
#     dataset_tools.start_convert()
    